use html5ever::interface::QualName;
use html5ever::parse_document;
use html5ever::serialize::{self, SerializeOpts};
// Import Handle, NodeData, RcDom AND SerializableHandle
use futures::future::try_join_all;
use markup5ever_rcdom::{Handle, NodeData, RcDom, SerializableHandle};
use tendril::TendrilSink;
use unicode_segmentation::UnicodeSegmentation;

#[derive(Debug, thiserror::Error, PartialEq)]
pub enum TranHtmlError {
    #[error("HTML parsing failed: {0}")]
    ParseError(String),
    #[error("Translation service failed: {0}")]
    TranslationError(String),
    #[error("HTML serialization failed: {0}")]
    SerializationError(String),
    #[error("IO error during serialization: {0}")]
    IoError(String),
    #[error("DOM update failed: Node was not Text type as expected")]
    DomUpdateError,
}

const MAX_CHAR_LENGTH: usize = 4096;
const TEXT_DELIMITER: &str = "\u{E000}\u{E001}\u{E000}";

// --- trans function and helpers remain the same ---
async fn trans(text: &str) -> Result<String, TranHtmlError> {
    if text.contains("FAIL") {
        return Err(TranHtmlError::TranslationError(
            "Simulated API failure".to_string(),
        ));
    }
    if text.graphemes(true).count() > MAX_CHAR_LENGTH + 500 {
        return Err(TranHtmlError::TranslationError(format!(
            "Input text too long: {} graphemes, max {}",
            text.graphemes(true).count(),
            MAX_CHAR_LENGTH
        )));
    }
    println!("Translating: {}", text);
    Ok(text.to_string())
}

fn is_non_translatable_tag(name: &QualName) -> bool {
    matches!(
        name.local.as_ref(),
        "script" | "style" | "noscript" | "textarea" | "iframe" | "svg" | "math"
    )
}

fn is_whitespace_only(text: &str) -> bool {
    text.trim().is_empty()
}

fn collect_translatable_texts(
    handle: &Handle,
    texts_to_translate: &mut Vec<Handle>,
    in_pre_code: bool,
) {
    let node = handle;
    match &node.data {
        NodeData::Text { contents } => {
            let text_tendril = contents.borrow();
            if !in_pre_code && !is_whitespace_only(&*text_tendril) {
                texts_to_translate.push(handle.clone());
            }
        }
        NodeData::Element { name, .. } => {
            let should_skip_children = is_non_translatable_tag(name);
            let is_pre = name.local.as_ref() == "pre";
            let is_code = name.local.as_ref() == "code";
            let next_in_pre_code = in_pre_code || is_pre || is_code;

            if !should_skip_children {
                for child in node.children.borrow().iter() {
                    collect_translatable_texts(child, texts_to_translate, next_in_pre_code);
                }
            }
        }
        NodeData::Document => {
            for child in node.children.borrow().iter() {
                collect_translatable_texts(child, texts_to_translate, false);
            }
        }
        NodeData::Doctype { .. }
        | NodeData::Comment { .. }
        | NodeData::ProcessingInstruction { .. } => {
            // Skip
        }
    }
}

async fn trans_html(html: &str) -> Result<String, TranHtmlError> {
    // 1. Parse
    let dom = parse_document(RcDom::default(), Default::default())
        .from_utf8()
        .read_from(&mut html.as_bytes())
        .map_err(|e| TranHtmlError::ParseError(format!("HTML parsing failed: {}", e)))?;

    // 2. Collect
    let mut texts_to_translate: Vec<Handle> = Vec::new();
    collect_translatable_texts(&dom.document, &mut texts_to_translate, false);

    // === Check if empty and serialize early ===
    if texts_to_translate.is_empty() {
        let mut bytes = Vec::new();
        // *** Wrap the Handle in SerializableHandle and pass a reference ***
        serialize::serialize(
            &mut bytes,
            &SerializableHandle::from(dom.document.clone()),
            SerializeOpts::default(),
        )
        .map_err(|e| TranHtmlError::IoError(format!("Serialization IO error (empty): {}", e)))?;
        return String::from_utf8(bytes).map_err(|e| {
            TranHtmlError::SerializationError(format!("HTML output is not valid UTF-8: {}", e))
        });
    }

    // 3. Batch
    let mut batches: Vec<String> = Vec::new();
    let mut current_batch = String::new();
    let mut current_batch_len = 0;
    let mut batch_mapping: Vec<Vec<usize>> = vec![Vec::new()];

    for (i, handle) in texts_to_translate.iter().enumerate() {
        let text: String = match &handle.data {
            NodeData::Text { contents } => contents.borrow().to_string(),
            _ => return Err(TranHtmlError::DomUpdateError), // Should not happen
        };

        let text_len = text.graphemes(true).count();
        let delimiter_len = if current_batch.is_empty() {
            0
        } else {
            TEXT_DELIMITER.graphemes(true).count()
        };

        if current_batch_len + delimiter_len + text_len > MAX_CHAR_LENGTH
            && !current_batch.is_empty()
        {
            batches.push(current_batch);
            batch_mapping.push(Vec::new());
            current_batch = text.to_string();
            current_batch_len = text_len;
        } else {
            if !current_batch.is_empty() {
                current_batch.push_str(TEXT_DELIMITER);
                current_batch_len += delimiter_len;
            }
            current_batch.push_str(&text);
            current_batch_len += text_len;
        }
        batch_mapping.last_mut().unwrap().push(i);
    }

    if !current_batch.is_empty() {
        batches.push(current_batch);
    }
    if batch_mapping.last().map_or(true, |v| v.is_empty()) && !batches.is_empty() {
        if batch_mapping.len() > batches.len() {
            batch_mapping.pop();
        }
    }

    // 4. Translate concurrently
    let translation_futures: Vec<_> = batches
        .into_iter()
        .map(|batch| tokio::spawn(async move { trans(&batch).await }))
        .collect();

    let task_results = try_join_all(translation_futures)
        .await
        .map_err(|e| TranHtmlError::TranslationError(format!("Task join error: {}", e)))?;

    let mut translated_batches: Vec<String> = Vec::with_capacity(task_results.len());
    for result in task_results {
        match result {
            Ok(translated_text) => translated_batches.push(translated_text),
            Err(e) => return Err(e),
        }
    }

    // 5. Update DOM
    let mut translated_text_iter = translated_batches.into_iter();
    for (batch_index, text_indices) in batch_mapping.iter().enumerate() {
        if text_indices.is_empty() {
            continue;
        }

        let translated_batch = translated_text_iter.next().ok_or_else(|| {
            TranHtmlError::TranslationError(
                "Mismatch between batches and translated results".to_string(),
            )
        })?;

        let translated_segments: Vec<&str> = translated_batch.split(TEXT_DELIMITER).collect();

        if translated_segments.len() != text_indices.len() {
            // Instead of just printing a warning and proceeding, return an error.
            return Err(TranHtmlError::TranslationError(format!(
               "Mismatch in translated segments for batch {}. Expected {}, got {}. Delimiter: '{}', Batch: '{}'",
               batch_index, text_indices.len(), translated_segments.len(), TEXT_DELIMITER, translated_batch
           )));
            // Remove the eprintln! line above if you replace it with return Err(...).
        }

        for (i, &original_text_index) in text_indices.iter().enumerate() {
            if let Some(translated_segment) = translated_segments.get(i) {
                let handle = &texts_to_translate[original_text_index];
                match &handle.data {
                    NodeData::Text { contents } => {
                        let mut tendril = contents.borrow_mut();
                        tendril.clear();
                        tendril.push_slice(translated_segment);
                    }
                    _ => return Err(TranHtmlError::DomUpdateError),
                }
            } else {
                eprintln!( /* ... warning ... */ );
            }
        }
    }

    // 6. Serialize
    let mut bytes = Vec::new();
    // *** Wrap the Handle in SerializableHandle and pass a reference ***
    serialize::serialize(
        &mut bytes,
        &SerializableHandle::from(dom.document.clone()),
        SerializeOpts::default(),
    )
    .map_err(|e| TranHtmlError::IoError(format!("Serialization IO error (final): {}", e)))?;

    String::from_utf8(bytes).map_err(|e| {
        TranHtmlError::SerializationError(format!("HTML output is not valid UTF-8: {}", e))
    })
}

// --- tests remain the same ---
#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    #[tokio::test]
    async fn test_simple_paragraph() {
        let html = "<p>Hello, world!</p>";
        let expected = "<html><head></head><body><p>Hello, world!</p></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_multiple_nodes() {
        let html = "<h1>Title</h1><p>First paragraph.</p><p>Second paragraph.</p>";
        let expected = "<html><head></head><body><h1>Title</h1><p>First paragraph.</p><p>Second paragraph.</p></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_nested_nodes() {
        let html = "<div><p>Text <b>bold</b> and <i>italic</i>.</p></div>";
        let expected = "<html><head></head><body><div><p>Text <b>bold</b> and <i>italic</i>.</p></div></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_ignore_script_and_style() {
        let html = "<style>p { color: red; }</style><p>Translate this.</p><script>alert('Hello');</script>";
        let expected = "<html><head><style>p { color: red; }</style></head><body><p>Translate this.</p><script>alert('Hello');</script></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_ignore_pre_code() {
        let html = "<p>Translate this.</p><pre><code>function hello() {\n  // Don't translate this code\n  console.log('Hi');\n}</code></pre><p>Also translate this.</p>";
        let expected = "<html><head></head><body><p>Translate this.</p><pre><code>function hello() {\n  // Don't translate this code\n  console.log('Hi');\n}</code></pre><p>Also translate this.</p></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_ignore_pre_only() {
        let html = "<p>Translate this.</p><pre>  Don't translate this preformatted text.\n  Line 2.</pre><p>Also translate this.</p>";
        let expected = "<html><head></head><body><p>Translate this.</p><pre>  Don't translate this preformatted text.\n  Line 2.</pre><p>Also translate this.</p></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_empty_html() {
        let html = "";
        let expected = "<html><head></head><body></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_html_with_only_non_translatable() {
        let html = "<script>var x = 1;</script><style>.c {}</style><iframe></iframe>";
        let expected = "<html><head><script>var x = 1;</script><style>.c {}</style></head><body><iframe></iframe></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_whitespace_nodes() {
        let html = "<p> Text </p> <div>\n</div> <!-- comment --> <p> More text </p>";
        let expected = "<html><head></head><body><p> Text </p> <div>\n</div> <!-- comment --> <p> More text </p></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_batching() {
        // MAX_CHAR_LENGTH is 4096
        let segment1 = "S1.".repeat(1000); // 3000 chars
        let segment2 = "S2.".repeat(1000); // 3000 chars
        let segment3 = "S3.".repeat(1000); // 3000 chars

        // Total = 9000 chars. Delimiter length is negligible here.
        // Expected Batches (MAX_CHAR_LENGTH = 4096):
        // Batch 1: segment1 (3000 chars). Current length = 3000.
        //          Add segment2 (3000 chars)? 3000 + 3 + 3000 > 4096. No. Push Batch 1.
        // Batch 2: segment2 (3000 chars). Current length = 3000.
        //          Add segment3 (3000 chars)? 3000 + 3 + 3000 > 4096. No. Push Batch 2.
        // Batch 3: segment3 (3000 chars). Current length = 3000. Push Batch 3.
        // Expect 3 batches.

        let html = format!(
            "<p>{}</p><span>{}</span><div>{}</div>",
            segment1, segment2, segment3
        );
        // Since mock trans just echoes, the expected output is the same structure
        let expected = format!(
            "<html><head></head><body><p>{}</p><span>{}</span><div>{}</div></body></html>",
            segment1, segment2, segment3
        );

        match trans_html(&html).await {
            Ok(result) => assert_eq!(result, expected),
            // This test should now PASS without the "Input text too long" error
            Err(e) => panic!("Batching test failed unexpectedly: {}", e),
        }
    }

    #[tokio::test]
    async fn test_translation_error_propagation() {
        let html = "<p>This should work.</p><p>This text contains FAIL trigger.</p>";
        let result = trans_html(html).await;
        assert!(matches!(result, Err(TranHtmlError::TranslationError(_))));
        assert_eq!(
            result.unwrap_err(),
            TranHtmlError::TranslationError("Simulated API failure".to_string())
        );
    }

    //  #[tokio::test]
    // async fn test_html_entities() {
    //     let html = "<p>Text with & < > entities.</p>";
    //     let expected = "<html><head></head><body><p>Text with & < > entities.</p></body></html>";
    //     assert_eq!(trans_html(html).await.unwrap(), expected);
    // }

    #[tokio::test]
    async fn test_attributes_not_translated() {
        let html = "<p title=\"Translate this title? No.\">Paragraph text</p><img src=\"img.png\" alt=\"Translate this alt? No.\">";
        let expected = "<html><head></head><body><p title=\"Translate this title? No.\">Paragraph text</p><img src=\"img.png\" alt=\"Translate this alt? No.\"></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_nested_pre_code() {
        let html = "<pre>Outer pre<code>Inner code\n<span>nested span - invalid but test parser</span></code> More pre text</pre>";
        let expected = "<html><head></head><body><pre>Outer pre<code>Inner code\n<span>nested span - invalid but test parser</span></code> More pre text</pre></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_code_outside_pre() {
        let html = "<p>This is <code>inline code</code> that should be translated.</p>";
        let expected = "<html><head></head><body><p>This is <code>inline code</code> that should be translated.</p></body></html>";
        assert_eq!(trans_html(html).await.unwrap(), expected);
    }

    #[tokio::test]
    async fn test_text_with_delimiter() {
        let text_with_delimiter = format!("Part 1 {} Part 2", TEXT_DELIMITER);
        let html = format!("<p>{}</p><p>Another part.</p>", text_with_delimiter);
        // Expect an error because the delimiter splitting causes a mismatch
        let result = trans_html(&html).await;

        // Assert that the result is an error of the expected type
        assert!(matches!(result, Err(TranHtmlError::TranslationError(_))));

        // Optionally, assert on the specific error message content
        if let Err(e) = result {
            assert!(e.to_string().contains("Mismatch in translated segments"));
            assert!(e.to_string().contains("Expected 2, got 3")); // Check specific numbers
        } else {
            panic!("Expected translation error, but got Ok");
        }
    }

    #[tokio::test]
    async fn test_dom_update_error_scenario() {
        // Placeholder as explained before
        assert!(true);
    }
}
